package org.apache.tinkerpop.gremlin.server;

import io.grpc.Server;
import io.grpc.ServerBuilder;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.channel.*;
import io.netty.channel.epoll.EpollEventLoopGroup;
import io.netty.channel.epoll.EpollServerSocketChannel;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.util.concurrent.GenericFutureListener;
import io.netty.util.internal.logging.InternalLoggerFactory;
import io.netty.util.internal.logging.Slf4JLoggerFactory;
import org.apache.commons.lang3.SystemUtils;
import org.apache.tinkerpop.gremlin.server.op.OpLoader;
import org.apache.tinkerpop.gremlin.server.util.GraphServerGremlinExecutor;
import org.apache.tinkerpop.gremlin.server.util.LifeCycleHook;
import org.apache.tinkerpop.gremlin.server.util.MetricManager;
import org.apache.tinkerpop.gremlin.server.util.ThreadFactoryUtil;
import org.apache.tinkerpop.gremlin.structure.Graph;
import org.janusgraph.core.JanusGraph;
import org.janusgraph.graphdb.grpc.JanusGraphContextHandler;
import org.janusgraph.graphdb.grpc.JanusGraphManagerServiceImpl;
import org.janusgraph.graphdb.grpc.schema.SchemaManagerImpl;
import org.janusgraph.graphdb.management.JanusGraphManager;
import org.janusgraph.graphdb.server.JanusGraphSettings;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.studiox.graph.common.exception.GraphException;
import org.studiox.graph.janusgraph.latest.JanusGraphLatest;

import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.UndeclaredThrowableException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.concurrent.*;

public class GraphGremlinServer {

  static {
    // hook slf4j up to netty internal logging
    InternalLoggerFactory.setDefaultFactory(Slf4JLoggerFactory.INSTANCE);
  }

  private static final String SERVER_THREAD_PREFIX = "gremlin-server-";

  private static final Logger logger = LoggerFactory.getLogger(GraphGremlinServer.class);

  private final Settings settings;

  private Channel ch;

  private CompletableFuture<Void> serverStopped = null;
  private CompletableFuture<GraphServerGremlinExecutor> serverStarted = null;

  private final EventLoopGroup bossGroup;
  private final EventLoopGroup workerGroup;
  private final ExecutorService gremlinExecutorService;

  private final GraphServerGremlinExecutor graphServerGremlinExecutor;

  private final boolean isEpollEnabled;

  private GraphChannelizer graphChannelizer;

  public GraphGremlinServer(final Settings settings) {
    this(settings, null, null);
  }

  public GraphGremlinServer(final Settings settings, Graph graph) {
    this(settings, null, graph);
  }

  public GraphGremlinServer(
      final Settings settings, final ExecutorService gremlinExecutorService, Graph graph) {
    this.settings = settings;

    provideDefaultForGremlinPoolSize(settings);
    this.isEpollEnabled = settings.useEpollEventLoop && SystemUtils.IS_OS_LINUX;
    if (settings.useEpollEventLoop && !SystemUtils.IS_OS_LINUX) {
      logger.warn("cannot use epoll in non-linux env, falling back to NIO");
    }

    Runtime.getRuntime()
        .addShutdownHook(new Thread(() -> this.stop().join(), SERVER_THREAD_PREFIX + "shutdown"));

    final ThreadFactory threadFactoryBoss = ThreadFactoryUtil.create("boss-%d");

    // if linux os use epoll else fallback to nio based eventloop
    // epoll helps in reducing GC and has better  performance
    // http://netty.io/wiki/native-transports.html
    if (isEpollEnabled) {
      bossGroup = new EpollEventLoopGroup(settings.threadPoolBoss, threadFactoryBoss);
    } else {
      bossGroup = new NioEventLoopGroup(settings.threadPoolBoss, threadFactoryBoss);
    }

    final ThreadFactory threadFactoryWorker = ThreadFactoryUtil.create("worker-%d");
    if (isEpollEnabled) {
      workerGroup = new EpollEventLoopGroup(settings.threadPoolWorker, threadFactoryWorker);
    } else {
      workerGroup = new NioEventLoopGroup(settings.threadPoolWorker, threadFactoryWorker);
    }

    // use the ExecutorService returned from ServerGremlinExecutor as it might be initialized there
    graphServerGremlinExecutor =
        new GraphServerGremlinExecutor(settings, gremlinExecutorService, workerGroup, graph);
    this.gremlinExecutorService = graphServerGremlinExecutor.getGremlinExecutorService();

    // initialize the OpLoader with configurations being passed to each OpProcessor implementation
    // loaded
    OpLoader.init(settings);
  }

  public void start() {
    Server grpcServer = createGrpcServer();

    CompletableFuture<Void> grpcServerFuture =
        CompletableFuture.runAsync(
            () -> {
              try {
                grpcServer.start();
              } catch (IOException e) {
                throw new IllegalStateException(e);
              }
            });

    CompletableFuture<Void> gremlinServerFuture =
        startGremlinServer().thenAcceptAsync(GraphGremlinServer::configure);

    CompletableFuture<Void> serverStarted =
        CompletableFuture.allOf(gremlinServerFuture, grpcServerFuture);
  }

  public synchronized CompletableFuture<GraphServerGremlinExecutor> startGremlinServer() {
    if (serverStarted != null) {
      // server already started - don't get it rolling again
      return serverStarted;
    }

    serverStarted = new CompletableFuture<>();
    final CompletableFuture<GraphServerGremlinExecutor> serverReadyFuture = serverStarted;
    try {
      final ServerBootstrap b = new ServerBootstrap();

      // when high value is reached then the channel becomes non-writable and stays like that until
      // the
      // low value is so that there is time to recover
      b.childOption(
          ChannelOption.WRITE_BUFFER_WATER_MARK,
          new WriteBufferWaterMark(
              settings.writeBufferLowWaterMark, settings.writeBufferHighWaterMark));
      b.childOption(ChannelOption.ALLOCATOR, PooledByteBufAllocator.DEFAULT);

      // fire off any lifecycle scripts that were provided by the user. hooks get initialized during
      // ServerGremlinExecutor initialization
      graphServerGremlinExecutor
          .getHooks()
          .forEach(
              hook -> {
                logger.info("Executing start up {}", LifeCycleHook.class.getSimpleName());
                try {
                  hook.onStartUp(new LifeCycleHook.Context(logger));
                } catch (UnsupportedOperationException uoe) {
                  // if the user doesn't implement onStartUp the scriptengine will throw
                  // this exception.  it can safely be ignored.
                }
              });

      graphChannelizer = createGraphChannelizer(settings);
      graphChannelizer.init(graphServerGremlinExecutor);
      b.group(bossGroup, workerGroup).childHandler(graphChannelizer);
      if (isEpollEnabled) {
        b.channel(EpollServerSocketChannel.class);
      } else {
        b.channel(NioServerSocketChannel.class);
      }

      // bind to host/port and wait for channel to be ready
      b.bind(settings.host, settings.port)
          .addListener(
              (ChannelFutureListener)
                  channelFuture -> {
                    if (channelFuture.isSuccess()) {
                      ch = channelFuture.channel();
                      logger.info(
                          "Gremlin Server configured with worker thread pool of {}, gremlin pool of {} and boss thread pool of {}.",
                          settings.threadPoolWorker,
                          settings.gremlinPool,
                          settings.threadPoolBoss);
                      logger.info("Channel started at port {}.", settings.port);
                      serverReadyFuture.complete(graphServerGremlinExecutor);
                    } else {
                      serverReadyFuture.completeExceptionally(
                          new IOException(
                              String.format(
                                  "Could not bind to %s and %s - perhaps something else is bound to that address.",
                                  settings.host, settings.port)));
                    }
                  });
    } catch (Exception ex) {
      logger.error("Gremlin Server Error", ex);
      serverReadyFuture.completeExceptionally(ex);
    }

    return serverStarted;
  }

  private static GraphChannelizer createGraphChannelizer(final Settings settings) {
    try {
      final Class clazz = Class.forName(settings.channelizer);
      final Object o = clazz.newInstance();

      final GraphChannelizer c = (GraphChannelizer) o;
      if (c.supportsIdleMonitor()) {
        logger.info(
            "idleConnectionTimeout was set to {} which resolves to {} seconds when configuring this value - this feature will be {}",
            settings.idleConnectionTimeout,
            settings.idleConnectionTimeout / 1000,
            settings.idleConnectionTimeout < 1000 ? "disabled" : "enabled");
        logger.info(
            "keepAliveInterval was set to {} which resolves to {} seconds when configuring this value - this feature will be {}",
            settings.keepAliveInterval,
            settings.keepAliveInterval / 1000,
            settings.keepAliveInterval < 1000 ? "disabled" : "enabled");
      }

      return c;
    } catch (ClassNotFoundException cnfe) {
      logger.error(
          "Could not find {} implementation defined by the 'graphChannelizer' setting as: {}",
          GraphChannelizer.class.getName(),
          settings.channelizer);
      throw new RuntimeException(cnfe);
    } catch (Exception ex) {
      logger.error(
          "Class defined by the 'graphChannelizer' setting as: {} could not be properly instantiated as a {}",
          settings.channelizer,
          GraphChannelizer.class.getName());
      throw new RuntimeException(ex);
    }
  }

  public synchronized CompletableFuture<Void> stop() {
    if (serverStopped != null) {
      // shutdown has started so don't fire it off again
      return serverStopped;
    }

    serverStopped = new CompletableFuture<>();
    final CountDownLatch servicesLeftToShutdown = new CountDownLatch(3);

    // release resources in the OpProcessors (e.g. kill sessions)
    OpLoader.getProcessors()
        .forEach(
            (key, value) -> {
              logger.info("Shutting down OpProcessor[{}]", key);
              try {
                value.close();
              } catch (Exception ex) {
                logger.warn(
                    "Shutdown will continue but, there was an error encountered while closing "
                        + key,
                    ex);
              }
            });

    // it's possible that a channel might not be initialized in the first place if bind() fails
    // because
    // of port conflict.  in that case, there's no need to wait for the channel to close.
    if (null == ch) {
      servicesLeftToShutdown.countDown();
    } else {
      ch.close().addListener(f -> servicesLeftToShutdown.countDown());
    }

    logger.info("Shutting down thread pools.");

    try {
      if (gremlinExecutorService != null) {
        gremlinExecutorService.shutdown();
      }
    } finally {
      logger.debug("Shutdown Gremlin thread pool.");
    }

    try {
      workerGroup
          .shutdownGracefully()
          .addListener((GenericFutureListener) f -> servicesLeftToShutdown.countDown());
    } finally {
      logger.debug("Shutdown Worker thread pool.");
    }
    try {
      bossGroup
          .shutdownGracefully()
          .addListener((GenericFutureListener) f -> servicesLeftToShutdown.countDown());
    } finally {
      logger.debug("Shutdown Boss thread pool.");
    }

    // channel is shutdown as are the thread pools - time to kill graphs as nothing else should be
    // acting on them
    new Thread(
            () -> {
              if (graphServerGremlinExecutor != null) {
                graphServerGremlinExecutor
                    .getHooks()
                    .forEach(
                        hook -> {
                          logger.info("Executing shutdown {}", LifeCycleHook.class.getSimpleName());
                          try {
                            hook.onShutDown(new LifeCycleHook.Context(logger));
                          } catch (UnsupportedOperationException
                              | UndeclaredThrowableException uoe) {
                            // if the user doesn't implement onShutDown the scriptengine will throw
                            // this exception.  it can safely be ignored.
                          }
                        });
              }

              try {
                if (gremlinExecutorService != null) {
                  if (!gremlinExecutorService.awaitTermination(30000, TimeUnit.MILLISECONDS)) {
                    logger.warn(
                        "Gremlin thread pool did not fully terminate - continuing with shutdown process");
                  }
                }
              } catch (InterruptedException ie) {
                logger.warn(
                    "Timeout waiting for Gremlin thread pool to shutdown - continuing with shutdown process.");
              }

              try {
                servicesLeftToShutdown.await(30000, TimeUnit.MILLISECONDS);
              } catch (InterruptedException ie) {
                logger.warn(
                    "Timeout waiting for boss/worker thread pools to shutdown - continuing with shutdown process.");
              }

              if (graphServerGremlinExecutor != null) {
                graphServerGremlinExecutor
                    .getGraphManager()
                    .getGraphNames()
                    .forEach(
                        gName -> {
                          logger.debug("Closing Graph instance [{}]", gName);
                          try {
                            graphServerGremlinExecutor.getGraphManager().getGraph(gName).close();
                          } catch (Exception ex) {
                            logger.warn(
                                String.format("Exception while closing Graph instance [%s]", gName),
                                ex);
                          } finally {
                            logger.info("Closed Graph instance [{}]", gName);
                          }
                        });
              }

              // kills reporter threads. this is a last bit of cleanup that can be done. typically,
              // the jvm is headed
              // for shutdown which would obviously kill the reporters, but when it isn't they just
              // keep reporting.
              // removing them all will silent them up and release the appropriate resources.
              MetricManager.INSTANCE.removeAllReporters();

              // removing all the metrics should allow Gremlin Server to clean up the metrics
              // instance so that it can be
              // started again in the same JVM without those metrics initialized which generates a
              // warning and won't
              // reset to start values
              MetricManager.INSTANCE.removeAllMetrics();

              logger.info("Gremlin Server - shutdown complete");
              serverStopped.complete(null);
            },
            SERVER_THREAD_PREFIX + "stop")
        .start();

    return serverStopped;
  }

  public GraphServerGremlinExecutor getGraphServerGremlinExecutor() {
    return graphServerGremlinExecutor;
  }

  private static void provideDefaultForGremlinPoolSize(final Settings settings) {
    if (settings.gremlinPool == 0) {
      settings.gremlinPool = Runtime.getRuntime().availableProcessors();
    }
  }

  private Server createGrpcServer() {
    JanusGraphSettings janusGraphSettings = (JanusGraphSettings) settings;
    GraphManager graphManager = getGraphServerGremlinExecutor().getGraphManager();
    JanusGraphContextHandler janusGraphContextHandler = new JanusGraphContextHandler(graphManager);
    return ServerBuilder.forPort(janusGraphSettings.getGrpcServer().getPort())
        .addService(new JanusGraphManagerServiceImpl(janusGraphContextHandler))
        .addService(new SchemaManagerImpl(janusGraphContextHandler))
        .build();
  }

  private static void configure(GraphServerGremlinExecutor graphServerGremlinExecutor) {
    GraphManager graphManager = graphServerGremlinExecutor.getGraphManager();
    ((JanusGraphManager) graphManager)
        .configureGremlinExecutor(graphServerGremlinExecutor.getGremlinExecutor());
  }

  public static JanusGraphSettings createJanusGraphSettings() {
    JanusGraphSettings janusGraphSettings = null;
    InputStream inputStream = null;
    try {
      inputStream =
          Thread.currentThread().getContextClassLoader().getResourceAsStream("gremlin-server.yaml");
      janusGraphSettings = JanusGraphSettings.read(inputStream);
    } catch (Throwable t) {
      logger.error(t.getMessage(), t);
    } finally {
      if (null != inputStream) {
        try {
          inputStream.close();
        } catch (IOException e) {
          logger.error(e.getMessage(), e);
        }
      }
    }

    janusGraphSettings
        .scriptEngines
        .get("gremlin-groovy")
        .plugins
        .put(
            "org.apache.tinkerpop.gremlin.jsr223.ScriptFileGremlinPlugin",
            Collections.singletonMap(
                "files",
                new ArrayList<>(
                    Arrays.asList(
                        Thread.currentThread()
                            .getContextClassLoader()
                            .getResource("graph.groovy")
                            .getFile()))));

    return janusGraphSettings;
  }

  public static void main(String[] args) throws GraphException {
    JanusGraphLatest janusGraphLatest = new JanusGraphLatest(logger);
    janusGraphLatest.open();
    janusGraphLatest.createVertex("uniq1", "type1", null);
    JanusGraph janusGraph = janusGraphLatest.getJanusGraph();

    JanusGraphSettings janusGraphSettings = GraphGremlinServer.createJanusGraphSettings();
    GraphGremlinServer graphGremlinServer = new GraphGremlinServer(janusGraphSettings, janusGraph);
    graphGremlinServer.start();
  }
}
